package model

import (
	"cluster-cache/common"
	"github.com/cespare/xxhash/v2"
	"sync"
)

type KeyValue interface {
	Query(key []byte) []byte
	Add(key []byte, value []byte)
	Delete(key []byte)
}

type MultiMap struct {
	store []map[string]string
	locks []*sync.RWMutex
	size  uint64
}

func NewMultiMap(size int) KeyValue {
	multiMap := &MultiMap{
		store: make([]map[string]string, size),
		size:  uint64(size),
	}
	multiMap.locks = make([]*sync.RWMutex, size)
	for i := 0; i < size; i++ {
		multiMap.locks[i] = new(sync.RWMutex)
		multiMap.store[i] = make(map[string]string)
	}
	return multiMap
}

func (m *MultiMap) OnOpen() error {
	return nil
}

func (m *MultiMap) Query(key []byte) []byte {
	bucket := xxhash.Sum64(key) % m.size
	ks := common.BytesToString(key)
	m.locks[bucket].RLock()
	if val, ok := m.store[bucket][ks]; ok {
		m.locks[bucket].RUnlock()
		return common.StringToBytes(val)
	} else {
		m.locks[bucket].RUnlock()
		return nil
	}
}

func (m *MultiMap) Add(key []byte, value []byte) {
	bucket := xxhash.Sum64(key) % m.size
	m.locks[bucket].Lock()
	m.store[bucket][string(key)] = string(value)
	m.locks[bucket].Unlock()
}

func (m *MultiMap) Delete(key []byte) {
	bucket := xxhash.Sum64(key) % m.size
	m.locks[bucket].Lock()
	delete(m.store[bucket], string(key))
	m.locks[bucket].Unlock()
}

func (m *MultiMap) Batch(keys [][]byte, values [][]byte) error {
	for i := 0; i < len(keys); i++ {
		bucket := xxhash.Sum64(keys[i]) % m.size
		m.locks[bucket].Lock()
		m.store[bucket][string(keys[i])] = string(values[i])
		m.locks[bucket].Unlock()
	}
	return nil
}
